"""
    生成需要输入大模型的Prompt, 方便多进程执行
"""

import json 
import argparse
import pandas as pd 



def prompt(input_table, output_json):
    df = pd.read_excel(input_table,dtype="str")
    title = df.columns.tolist() 

    print(input_table)

    cluster_group = df.groupby("cluster_name")
    result_json = []
    for cluster_id, cluster_df in cluster_group:
        # 数基因类别的数量 
        gene_group = cluster_df.groupby("gene_description")
        gene_result = []
        for gene, gene_df in gene_group:
            gene_result.append([gene, gene_df.shape[0]])
        gene_df = pd.DataFrame(gene_result, columns=["gene_description", "count"])
        gene_df = gene_df.sort_values(by="count", ascending=False)
        gene_prompt = gene_df.values.tolist() 
        
        # pfam的提示词
        pfams = [str(x) for x in cluster_df["pfam"].tolist()]
        pfam_dict = {}
        for pfam in pfams:
            if pfam != "nan":
                pfam = eval(pfam)
                if isinstance(pfam, list):
                    for p in pfam:
                        if p not in pfam_dict:
                            pfam_dict[p] = 0
                        pfam_dict[p] += 1             
        pfam_list = []
        for k,v in pfam_dict.items():
            pfam_list.append([k,v])
        pfam_df = pd.DataFrame(pfam_list, columns=["pfam","count"])
        pfam_df = pfam_df.sort_values(by="count", ascending=False)
        pfam_prompt = pfam_df.values.tolist()
        
        # gog_category的提示词
        gog_group = cluster_df.groupby("gog_category")
        gog_result = []
        for gog, gog_df in gog_group:
            gog_result.append([gog, gog_df.shape[0]])
        gog_df = pd.DataFrame(gog_result, columns=["gog_category", "count"])
        gog_df = gog_df.sort_values(by="count", ascending=False)
        gog_prompt = gog_df.values.tolist()
        
        result_json.append({
            cluster_id: {
                "gene_prompt": gene_prompt,
                "pfam_prompt": pfam_prompt,
                "gog_prompt": gog_prompt
            }
        })
        
    with open(output_json, "w") as f:
        json.dump(result_json, f, indent=4)
        
        
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-input","--input")
    parser.add_argument("-output","--output")
    args = parser.parse_args()
    
    return args 


if __name__ == "__main__":
    params = parse_args()
    input_table = params.input
    output_json = params.output
    prompt(input_table, output_json)    